﻿namespace Bioinformatics

open DiffLCS

open System

// Integer-based
type DNABase = 
    | A = 0
    | C = 1
    | G = 2
    | T = 3

type FASTA() =
    // Could be more efficient
    static member public Parse(a : string) =
        a.Split([| "\r\n"; "\n" |], StringSplitOptions.RemoveEmptyEntries)
            |> Seq.filter (fun line -> String.length line <> 0 && line.[0] <> '>' && line.[0] <> ';')
            |> Seq.map (fun line -> line.Replace(" ", ""))
            |> Seq.concat
            |> Seq.toList
            |> List.map (fun c -> DNABase.TryParse(c.ToString()))
            |> List.filter (fun (succes, res) -> succes)
            |> List.map (fun (succes, res : DNABase) -> res)
            |> List.toArray

type LocalAlignment(a : DNABase[], b : DNABase[], subcost : int[,], gapcost : int) =
    // TODO: Share code with GlobalAlignment
    do
        if subcost.GetLength(0) <> 4 || subcost.GetLength(1) <> 4 then
            raise (new ArgumentException("The subcost matrix is not suitable."))

    let costMatrix = LocalAlignment.costMatrixSeq(a, b, subcost, gapcost)

    /// Returns a Local Alignment Recursion backtrack matrix. This is done sequential.
    static member private costMatrixSeq (a : DNABase[], b : DNABase[], subcost : int[,], gapcost : int) =
        let (n, m) = (a.Length, b.Length)
        let C = Array2D.zeroCreate (n + 1) (m + 1)

        // Array does not need to be initialized with the gapcost 
        // along the first row and first column if it is negative
        // because then 0 (local alignment) is higher
        if gapcost > 0 then
            for i = 1 to n do
                C.[i, 0] <- gapcost * i

            for j = 1 to m do
                C.[0, j] <- gapcost * j
    
        for i = 1 to n do
            for j = 1 to m do
                // TODO: Benchmark the performance of "max" instead of "Math.Max"
                let q = C.[i - 1, j - 1] + subcost.[int a.[i - 1], int b.[j - 1]]
                let p = C.[i - 1, j] + gapcost
                let r = C.[i, j - 1] + gapcost
                C.[i, j] <- max (max q p)
                                (max r 0)

        C

    /// Returns a copy of the cost/backtrack-matrix
    member public this.CostMatrix with get() = Array2D.copy costMatrix

    // TODO: Do this when calculating the costmatrix
    member private this.findOptimalEntry() =
        let maxEntryIndex = ref (0, 0)
        let maxEntryValue = ref Int32.MinValue

        // INFO: Maybe not the best optimal alignment
        Array2D.iteri (fun i j (a : int) -> if a > !maxEntryValue then
                                                maxEntryIndex := i, j
                                                maxEntryValue := a) costMatrix
        
        let i, j = !maxEntryIndex

        (i, j, !maxEntryValue)

    /// Returns the trace of an optimal alignment.
    member public this.OptimalTrace() =
        let C = costMatrix

        let rec visit i j =
            if C.[i, j] = 0 then
                []
            elif i = 0 then
                (-1, int b.[j - 1]) :: visit i (j - 1)
            elif j = 0 then
                (int a.[i - 1], -1) :: visit (i - 1) j
            else
                let q = C.[i - 1, j - 1] + subcost.[int a.[i - 1], int b.[j - 1]]
                let p = C.[i - 1, j] + gapcost
                let r = C.[i, j - 1] + gapcost

                let maxcost0 = max q (max p r)

                if maxcost0 = q then // Match/Substitution
                    (int a.[i - 1], int b.[j - 1]) :: visit (i - 1) (j - 1)
                elif maxcost0 = p then // Deletion
                    (int a.[i - 1], -1) :: visit (i - 1) j
                else // Insertion
                    (-1, int b.[j - 1]) :: visit i (j - 1)

        let (i, j, _) = this.findOptimalEntry()

        visit i j |> List.rev

    member public this.OptimalCost () =
        let _, _, optimalCost = this.findOptimalEntry()

        optimalCost

    member public this.Optimal () =
        let optimalalignment = this.OptimalTrace()
       
        let top = List.map fst optimalalignment
        let bottom = List.map snd optimalalignment

        let mapping = function
                        | -1 -> '-'
                        | x -> enum<DNABase>(x).ToString().[0]
        let f xs = xs |> List.map mapping |> List.toString

        (f top, f bottom)

type GlobalAlignment(a : DNABase[], b : DNABase[], subcost : int[,], gapcost : int) =
    do
        if subcost.GetLength(0) <> 4 || subcost.GetLength(1) <> 4 then
            raise (new ArgumentException("The subcost matrix is not suitable."))

    let costMatrix = GlobalAlignment.costMatrixSeq(a, b, subcost, gapcost)

    // TODO: Should it be properties or methods?

    /// Returns a Global Alignment Recursion backtrack matrix. This is done sequential.
    static member private costMatrixSeq (a : DNABase[], b : DNABase[], subcost : int[,], gapcost : int) =
//        let (n, m) = (a.Length, b.Length)
//        let C = Array2D.zeroCreate (n + 1) (m + 1)
//
//        for i = 1 to n do
//            C.[i, 0] <- if i % 3 = 0 then
//                            gapcost * (i / 3)
//                        else
//                            Int32.MinValue
//
//        for j = 1 to m do
//            C.[0, j] <- if j % 3 = 0 then
//                            gapcost * (j / 3)
//                        else
//                            Int32.MinValue
//    
//        for i = 1 to n do
//            for j = 1 to m do
//                let q = C.[i - 1, j - 1] + subcost.[int a.[i - 1], int b.[j - 1]]
//                let p = if i >= 3 then
//                            C.[i - 3, j] + gapcost
//                        else
//                            Int32.MinValue
//                let r = if j >= 3 then
//                            C.[i, j - 3] + gapcost
//                        else
//                            Int32.MinValue
//
//                C.[i, j] <- max q (max p r)
//
//        C
        let (n, m) = (a.Length, b.Length)
        let C = Array2D.zeroCreate (n + 1) (m + 1)

        for i = 1 to n do
            C.[i, 0] <- gapcost * i

        for j = 1 to m do
            C.[0, j] <- gapcost * j
    
        for i = 1 to n do
            for j = 1 to m do
                // TODO: Benchmark the performance of "max" instead of "Math.Max"
                C.[i, j] <- max (C.[i - 1, j - 1] + subcost.[int a.[i - 1], int b.[j - 1]])
                                (max (C.[i - 1, j] + gapcost) (C.[i, j - 1] + gapcost))

        C

    /// Returns a copy of the cost/backtrack-matrix
    member public this.CostMatrix with get() = Array2D.copy costMatrix

    /// Returns the trace of an optimal alignment.
    member public this.OptimalTrace() =
//        let C = costMatrix
//
//        let rec visit i j =
//            if i = 0 && j = 0 then
//                []
//            elif i = 0 then
//                (-1, int b.[j - 1]) :: visit i (j - 1)
//            elif j = 0 then
//                (int a.[i - 1], -1) :: visit (i - 1) j
//            else
//                let q = C.[i - 1, j - 1] + subcost.[int a.[i - 1], int b.[j - 1]]
//                let p = if i >= 3 then
//                            C.[i - 3, j] + gapcost
//                        else
//                            Int32.MinValue
//                let r = if j >= 3 then
//                            C.[i, j - 3] + gapcost
//                        else
//                            Int32.MinValue
//
//                let maxcost0 = max q (max p r)
//
//                if maxcost0 = q then // Match/Substitution
//                    (int a.[i - 1], int b.[j - 1]) :: visit (i - 1) (j - 1)
//                elif maxcost0 = p then // Deletion
//                    (int a.[i - 1], -1) :: (int a.[i - 2], -1) :: (int a.[i - 3], -1) :: visit (i - 3) j
//                else // Insertion
//                    (-1, int b.[j - 1]) :: (-1, int b.[j - 2]) :: (-1, int b.[j - 2]) :: visit i (j - 3)
//
//        visit a.Length b.Length |> List.rev
        let C = costMatrix

        let rec visit i j =
            if i = 0 && j = 0 then
                []
            elif i = 0 then
                (-1, int b.[j - 1]) :: visit i (j - 1)
            elif j = 0 then
                (int a.[i - 1], -1) :: visit (i - 1) j
            else
                let q = C.[i - 1, j - 1] + subcost.[int a.[i - 1], int b.[j - 1]]
                let p = C.[i - 1, j] + gapcost
                let r = C.[i, j - 1] + gapcost

                let maxcost0 = max q (max p r)

                if maxcost0 = q then // Match/Substitution
                    (int a.[i - 1], int b.[j - 1]) :: visit (i - 1) (j - 1)
                elif maxcost0 = p then // Deletion
                    (int a.[i - 1], -1) :: visit (i - 1) j
                else // Insertion
                    (-1, int b.[j - 1]) :: visit i (j - 1)

        visit a.Length b.Length |> List.rev

    /// Returns the count of optimal alignments.
    member public this.OptimalCount () =
//        let C = costMatrix
//
//        let rec visit i j =
//            if i = 0 && j = 0 then
//                1
//            elif i = 0 then
//                1 + visit i (j - 1)
//            elif j = 0 then
//                1 + visit (i - 1) j
//            else
//                let q = C.[i - 1, j - 1] + subcost.[int a.[i - 1], int b.[j - 1]]
//                let p = if i >= 3 then
//                            C.[i - 3, j] + gapcost
//                        else
//                            Int32.MinValue
//                let r = if j >= 3 then
//                            C.[i, j - 3] + gapcost
//                        else
//                            Int32.MinValue
//
//                let maxcost0 = max q (max p r)
//                let mutable count0 = 0
//
//                if maxcost0 = q then // Match/Substitution
//                    count0 <- count0 + visit (i - 1) (j - 1)
//
//                if maxcost0 = p then // Deletion
//                    count0 <- count0 + visit (i - 3) j
//
//                if maxcost0 = r then // Insertion
//                    count0 <- count0 + visit i (j - 3)
//
//                count0
//
//        visit a.Length b.Length
        let C = costMatrix

        let rec visit i j =
            if i = 0 && j = 0 then
                1
            elif i = 0 then
                1 + visit i (j - 1)
            elif j = 0 then
                1 + visit (i - 1) j
            else
                let q = C.[i - 1, j - 1] + subcost.[int a.[i - 1], int b.[j - 1]]
                let p = C.[i - 1, j] + gapcost
                let r = C.[i, j - 1] + gapcost

                let maxcost0 = max q (max p r)
                let mutable count0 = 0

                if maxcost0 = q then // Match/Substitution
                    count0 <- count0 + visit (i - 1) (j - 1)

                if maxcost0 = p then // Deletion
                    count0 <- count0 + visit (i - 1) j

                if maxcost0 = r then // Insertion
                    count0 <- count0 + visit i (j - 1)

                count0

        visit a.Length b.Length

    member public this.OptimalCost () =
        costMatrix.[a.Length, b.Length]

    member public this.Optimal () =
        let optimalalignment = this.OptimalTrace()
       
        let top = List.map fst optimalalignment
        let bottom = List.map snd optimalalignment

        let mapping = function
                        | -1 -> '-'
                        | x -> enum<DNABase>(x).ToString().[0]
        let f xs = xs |> List.map mapping |> List.toString

        (f top, f bottom)